{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "from torchvision import datasets, transforms\n",
    "import argparse\n",
    "import time\n",
    "import os\n",
    "from torch.autograd import Variable\n",
    "\n",
    "#setup training parameters\n",
    "parser = argparse.ArgumentParser(description='PyTorch MNIST Training')\n",
    "parser.add_argument('--batch-size', type=int, default=128, metavar='N',\n",
    "                    help='input batch size for training (default: 128)')\n",
    "parser.add_argument('--test-batch-size', type=int, default=128, metavar='N',\n",
    "                    help='input batch size for testing (default: 128)')\n",
    "parser.add_argument('--lr', type=float, default=0.01, metavar='LR',\n",
    "                    help='learning rate')\n",
    "parser.add_argument('--no-cuda', action='store_true', default=False,\n",
    "                    help='disables CUDA training')\n",
    "parser.add_argument('--seed', type=int, default=1, metavar='S',\n",
    "                    help='random seed (default: 1)')\n",
    "parser.add_argument('--model-dir', default='./model-mnist-cnn',\n",
    "                    help='directory of model for saving checkpoint')\n",
    "parser.add_argument('--random', default=True,\n",
    "                    help='random initialization for PGD')\n",
    "\n",
    "\n",
    "#FGSM: num-steps:1 step-size:0.031   PGD-20: num-steps:20 step-size:0.003  \n",
    "parser.add_argument('--epsilon', default=0.031,\n",
    "                    help='perturbation')\n",
    "parser.add_argument('--num-steps', default=1,\n",
    "                    help='perturb number of steps, FGSM: 1, PGD-20: 20')\n",
    "parser.add_argument('--step-size', default=0.031,\n",
    "                    help='perturb step size, FGSM: 0.031, PGD-20: 0.003')\n",
    "\n",
    "args = parser.parse_args(args=[]) \n",
    "\n",
    "if not os.path.exists(args.model_dir):\n",
    "    os.makedirs(args.model_dir)\n",
    "        \n",
    "# Judge cuda is available or not\n",
    "use_cuda = not args.no_cuda and torch.cuda.is_available()\n",
    "#device = torch.device(\"cuda\" if use_cuda else \"cpu\")\n",
    "device = torch.device(\"cpu\")\n",
    "\n",
    "torch.manual_seed(args.seed)\n",
    "kwargs = {'num_workers': 1, 'pin_memory': True} if use_cuda else {}\n",
    "\n",
    "# Setup data loader\n",
    "transform=transforms.Compose([\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize((0.1307,), (0.3081,))\n",
    "        ])\n",
    "trainset = datasets.MNIST('../data', train=True, download=True,\n",
    "                   transform=transform)\n",
    "testset = datasets.MNIST('../data', train=False,\n",
    "                   transform=transform)\n",
    "train_loader = torch.utils.data.DataLoader(trainset,batch_size=args.batch_size, shuffle=True,**kwargs)\n",
    "test_loader = torch.utils.data.DataLoader(testset,batch_size=args.test_batch_size, shuffle=False, **kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define CNN\n",
    "class Net(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Net, self).__init__()\n",
    "        # in_channels:1  out_channels:32  kernel_size:3  stride:1\n",
    "        self.conv1 = nn.Conv2d(1, 32, 3, 1)\n",
    "        # in_channels:32  out_channels:64  kernel_size:3  stride:1\n",
    "        self.conv2 = nn.Conv2d(32, 64, 3, 1)\n",
    "        self.fc1 = nn.Linear(9216, 128)\n",
    "        self.fc2 = nn.Linear(128, 10)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.conv1(x)\n",
    "        x = F.relu(x)\n",
    "        x = self.conv2(x)\n",
    "        x = F.relu(x)\n",
    "        x = F.max_pool2d(x, 2)\n",
    "        x = torch.flatten(x, 1)\n",
    "        x = self.fc1(x)\n",
    "        x = F.relu(x)\n",
    "        x = self.fc2(x)\n",
    "        output = F.log_softmax(x, dim=1)\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _pgd_whitebox(model,\n",
    "                  X,\n",
    "                  y,\n",
    "                  epsilon=args.epsilon,\n",
    "                  num_steps=args.num_steps,\n",
    "                  step_size=args.step_size):\n",
    "    out = model(X)\n",
    "    err = (out.data.max(1)[1] != y.data).float().sum()\n",
    "    X_pgd = Variable(X.data, requires_grad=True)\n",
    "    if args.random:\n",
    "        random_noise = torch.FloatTensor(*X_pgd.shape).uniform_(-epsilon, epsilon).to(device)\n",
    "        X_pgd = Variable(X_pgd.data + random_noise, requires_grad=True)\n",
    "\n",
    "    for _ in range(num_steps):\n",
    "        opt = optim.SGD([X_pgd], lr=1e-3)\n",
    "        opt.zero_grad()\n",
    "\n",
    "        with torch.enable_grad():\n",
    "            loss = nn.CrossEntropyLoss()(model(X_pgd), y)\n",
    "        loss.backward()\n",
    "        eta = step_size * X_pgd.grad.data.sign()\n",
    "        X_pgd = Variable(X_pgd.data + eta, requires_grad=True)\n",
    "        eta = torch.clamp(X_pgd.data - X.data, -epsilon, epsilon)\n",
    "        X_pgd = Variable(X.data + eta, requires_grad=True)\n",
    "        X_pgd = Variable(torch.clamp(X_pgd, 0, 1.0), requires_grad=True)\n",
    "    err_pgd = (model(X_pgd).data.max(1)[1] != y.data).float().sum()\n",
    "    return err, err_pgd\n",
    "\n",
    "def eval_adv_test_whitebox(model, device, test_loader):\n",
    "    #evaluate model by white-box attack\n",
    "    model.eval()\n",
    "    robust_err_total = 0\n",
    "    natural_err_total = 0\n",
    "\n",
    "    for data, target in test_loader:\n",
    "        data, target = data.to(device), target.to(device)\n",
    "        # pgd attack\n",
    "        X, y = Variable(data, requires_grad=True), Variable(target)\n",
    "        err_natural, err_robust = _pgd_whitebox(model, X, y)\n",
    "        robust_err_total += err_robust\n",
    "        natural_err_total += err_natural\n",
    "    print('natural_accuracy: {:.2f}%'.format(0.01 * (10000-natural_err_total)))\n",
    "    print('robust_accuracy: : {:.2f}%'.format(0.01 * (10000-robust_err_total)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "natural_accuracy: 96.89%\n",
      "robust_accuracy: : 92.64%\n"
     ]
    }
   ],
   "source": [
    "def main():\n",
    "    model = Net().to(device)\n",
    "    model.load_state_dict(torch.load(os.path.join(args.model_dir, 'final_model.pt')))\n",
    "    eval_adv_test_whitebox(model, device, test_loader)\n",
    "if __name__ == '__main__':\n",
    "    main()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
